//GhidraMCPServer.java
//@category MCP
//@description Ghidra MCP 后端服务器 - 为MCP集成提供REST API
//@author 卡耐基的反复手
//@toolbar 

import com.sun.net.httpserver.*;
import com.google.gson.*;
import ghidra.app.script.GhidraScript;
import ghidra.app.decompiler.*;
import ghidra.program.model.listing.*;
import ghidra.program.model.symbol.*;
import ghidra.program.model.address.*;
import ghidra.program.model.mem.Memory;
import ghidra.program.model.mem.MemoryBlock;
import ghidra.program.model.data.DataType;
import ghidra.util.task.TaskMonitor;   

import java.io.*;
import java.net.InetSocketAddress;
import java.util.*;

public class GhidraMCPServer extends GhidraScript {
    
    private HttpServer server;
    private Gson gson = new GsonBuilder().setPrettyPrinting().create();
    private DecompInterface decompiler;
    private volatile boolean running = true;
    
    @Override
    public void run() throws Exception {
        // 初始化反编译器
        decompiler = new DecompInterface();
        decompiler.openProgram(currentProgram);
        
        // 创建HTTP服务器
        int port =  2233; //askInt("服务器端口", "请输入端口号:");
        server = HttpServer.create(new InetSocketAddress(port), 0);
        
        // 注册API端点
        registerEndpoints();
        
        // 启动服务器
        server.setExecutor(null);
        server.start();
        
        println("========================================");
        println("Ghidra MCP 服务器启动成功！");
        println("端口: " + port);
        println("程序: " + currentProgram.getName());
        println("========================================");
        println("");
        println("服务器正在运行中...");
        println("要停止服务器，请:");
        println("1. 在Ghidra Script Manager中右键点击此脚本");
        println("2. 选择 'Cancel' 或关闭Ghidra");
        println("");
        println("API端点:");
        println("  健康检查: http://localhost:" + port + "/health");
        println("  程序信息: http://localhost:" + port + "/api/info");
        println("  函数列表: http://localhost:" + port + "/api/functions");
        println("========================================");
        
        // 保持脚本运行
        try {
            while (running && !monitor.isCancelled()) {
                Thread.sleep(1000);
            }
        } catch (InterruptedException e) {
            // 正常中断
        } finally {
            stopServer();
        }
    }
    
    private void registerEndpoints() {
        server.createContext("/api/decompile", this::handleDecompile);
        server.createContext("/api/function", this::handleFunction);
        server.createContext("/api/functions", this::handleListFunctions);
        server.createContext("/api/strings", this::handleStrings);
        server.createContext("/api/xrefs", this::handleXRefs);
        server.createContext("/api/calls", this::handleCalls);
        server.createContext("/api/rename", this::handleRename);
        server.createContext("/api/comment", this::handleComment);
        server.createContext("/api/info", this::handleProgramInfo);
        server.createContext("/api/asm", this::handleAsm);
        server.createContext("/api/memory", this::handleMemory);
        server.createContext("/api/memory/regions", this::handleMemoryRegions);
        server.createContext("/api/memory/search", this::handleMemorySearch);
        server.createContext("/api/memory/write", this::handleMemoryWrite);
        server.createContext("/health", this::handleHealth);
    }
    
    // ========== API处理器 ==========
    
    private void handleDecompile(HttpExchange exchange) throws IOException {
        if (!"POST".equals(exchange.getRequestMethod())) {
            sendError(exchange, 405, "不允许的方法");
            return;
        }
        
        try {
            JsonObject request = parseRequest(exchange);
            String addressStr = request.get("address").getAsString();
            Address addr = getAddress(addressStr);
            
            Function func = currentProgram.getFunctionManager().getFunctionContaining(addr);
            if (func == null) {
                sendError(exchange, 404, "地址处没有函数: " + addressStr);
                return;
            }
            
            // 反编译
            DecompileResults results = decompiler.decompileFunction(func, 30, monitor);
            
            if (!results.decompileCompleted()) {
                sendError(exchange, 500, "反编译失败: " + results.getErrorMessage());
                return;
            }
            
            JsonObject response = new JsonObject();
            response.addProperty("address", func.getEntryPoint().toString());
            response.addProperty("name", func.getName());
            response.addProperty("signature", func.getSignature().getPrototypeString());
            response.addProperty("decompiled_code", results.getDecompiledFunction().getC());
            
            sendResponse(exchange, 200, response);
            
        } catch (Exception e) {
            sendError(exchange, 500, "错误: " + e.getMessage());
        }
    }
    
    private void handleFunction(HttpExchange exchange) throws IOException {
        try {
            JsonObject request = parseRequest(exchange);
            String addressStr = request.get("address").getAsString();
            Address addr = getAddress(addressStr);
            
            Function func = currentProgram.getFunctionManager().getFunctionContaining(addr);
            if (func == null) {
                sendError(exchange, 404, "地址处没有函数: " + addressStr);
                return;
            }
            
            JsonObject response = new JsonObject();
            response.addProperty("address", func.getEntryPoint().toString());
            response.addProperty("name", func.getName());
            response.addProperty("signature", func.getSignature().getPrototypeString());
            response.addProperty("body_size", func.getBody().getNumAddresses());
            response.addProperty("is_thunk", func.isThunk());
            
            // 参数信息
            JsonArray params = new JsonArray();
            for (Parameter param : func.getParameters()) {
                JsonObject paramObj = new JsonObject();
                paramObj.addProperty("name", param.getName());
                paramObj.addProperty("type", param.getDataType().getName());
                params.add(paramObj);
            }
            response.add("parameters", params);
            
            sendResponse(exchange, 200, response);
            
        } catch (Exception e) {
            sendError(exchange, 500, "错误: " + e.getMessage());
        }
    }
    
    private void handleListFunctions(HttpExchange exchange) throws IOException {
        try {
            JsonArray functions = new JsonArray();
            FunctionIterator funcIter = currentProgram.getFunctionManager().getFunctions(true);
            
            int limit = 100;
            int count = 0;
            
            while (funcIter.hasNext() && count < limit) {
                Function func = funcIter.next();
                JsonObject funcObj = new JsonObject();
                funcObj.addProperty("address", func.getEntryPoint().toString());
                funcObj.addProperty("name", func.getName());
                funcObj.addProperty("signature", func.getSignature().getPrototypeString());
                functions.add(funcObj);
                count++;
            }
            
            JsonObject response = new JsonObject();
            response.add("functions", functions);
            response.addProperty("count", count);
            response.addProperty("total", currentProgram.getFunctionManager().getFunctionCount());
            
            sendResponse(exchange, 200, response);
            
        } catch (Exception e) {
            sendError(exchange, 500, "错误: " + e.getMessage());
        }
    }
    
    private void handleStrings(HttpExchange exchange) throws IOException {
        try {
            JsonArray strings = new JsonArray();
            Listing listing = currentProgram.getListing();
            DataIterator dataIter = listing.getDefinedData(true);
            
            int count = 0;
            int limit = 500;
            
            while (dataIter.hasNext() && count < limit) {
                Data data = dataIter.next();
                if (data.hasStringValue()) {
                    JsonObject strObj = new JsonObject();
                    strObj.addProperty("address", data.getAddress().toString());
                    strObj.addProperty("value", data.getDefaultValueRepresentation());
                    strObj.addProperty("length", data.getLength());
                    strings.add(strObj);
                    count++;
                }
            }
            
            JsonObject response = new JsonObject();
            response.add("strings", strings);
            response.addProperty("count", count);
            
            sendResponse(exchange, 200, response);
            
        } catch (Exception e) {
            sendError(exchange, 500, "错误: " + e.getMessage());
        }
    }
    
    private void handleXRefs(HttpExchange exchange) throws IOException {
        try {
            JsonObject request = parseRequest(exchange);
            String addressStr = request.get("address").getAsString();
            Address addr = getAddress(addressStr);
            
            JsonArray xrefs = new JsonArray();
            ReferenceIterator refIter = currentProgram.getReferenceManager().getReferencesTo(addr);
            
            int count = 0;
            while (refIter.hasNext()) {
                Reference ref = refIter.next();
                JsonObject refObj = new JsonObject();
                refObj.addProperty("from", ref.getFromAddress().toString());
                refObj.addProperty("type", ref.getReferenceType().getName());
                xrefs.add(refObj);
                count++;
            }
            
            JsonObject response = new JsonObject();
            response.add("xrefs", xrefs);
            response.addProperty("count", count);
            
            sendResponse(exchange, 200, response);
            
        } catch (Exception e) {
            sendError(exchange, 500, "错误: " + e.getMessage());
        }
    }
    
    private void handleCalls(HttpExchange exchange) throws IOException {
        try {
            JsonObject request = parseRequest(exchange);
            String addressStr = request.get("address").getAsString();
            Address addr = getAddress(addressStr);
            
            Function func = currentProgram.getFunctionManager().getFunctionContaining(addr);
            if (func == null) {
                sendError(exchange, 404, "地址处没有函数: " + addressStr);
                return;
            }
            
            JsonArray calls = new JsonArray();
            Set<Function> calledFunctions = func.getCalledFunctions(monitor);
            
            for (Function called : calledFunctions) {
                JsonObject callObj = new JsonObject();
                callObj.addProperty("address", called.getEntryPoint().toString());
                callObj.addProperty("name", called.getName());
                calls.add(callObj);
            }
            
            JsonObject response = new JsonObject();
            response.add("called_functions", calls);
            response.addProperty("count", calledFunctions.size());
            
            sendResponse(exchange, 200, response);
            
        } catch (Exception e) {
            sendError(exchange, 500, "错误: " + e.getMessage());
        }
    }
    
    private void handleRename(HttpExchange exchange) throws IOException {
        if (!"POST".equals(exchange.getRequestMethod())) {
            sendError(exchange, 405, "不允许的方法");
            return;
        }
        
        try {
            JsonObject request = parseRequest(exchange);
            String addressStr = request.get("address").getAsString();
            String newName = request.get("name").getAsString();
            Address addr = getAddress(addressStr);
            
            Function func = currentProgram.getFunctionManager().getFunctionContaining(addr);
            if (func == null) {
                sendError(exchange, 404, "地址处没有函数: " + addressStr);
                return;
            }
            
            int txId = currentProgram.startTransaction("重命名函数");
            try {
                func.setName(newName, SourceType.USER_DEFINED);
                currentProgram.endTransaction(txId, true);
            } catch (Exception e) {
                currentProgram.endTransaction(txId, false);
                throw e;
            }
            
            JsonObject response = new JsonObject();
            response.addProperty("success", true);
            response.addProperty("address", addressStr);
            response.addProperty("new_name", newName);
            
            sendResponse(exchange, 200, response);
            
        } catch (Exception e) {
            sendError(exchange, 500, "错误: " + e.getMessage());
        }
    }
    
    private void handleComment(HttpExchange exchange) throws IOException {
        if (!"POST".equals(exchange.getRequestMethod())) {
            sendError(exchange, 405, "不允许的方法");
            return;
        }
        
        try {
            JsonObject request = parseRequest(exchange);
            String addressStr = request.get("address").getAsString();
            String comment = request.get("comment").getAsString();
            Address addr = getAddress(addressStr);
            
            int txId = currentProgram.startTransaction("添加注释");
            try {
                Listing listing = currentProgram.getListing();
                // 使用新的方法替代过时的setComment
                listing.setComment(addr, CodeUnit.EOL_COMMENT, comment);
                currentProgram.endTransaction(txId, true);
            } catch (Exception e) {
                currentProgram.endTransaction(txId, false);
                throw e;
            }
            
            JsonObject response = new JsonObject();
            response.addProperty("success", true);
            response.addProperty("address", addressStr);
            
            sendResponse(exchange, 200, response);
            
        } catch (Exception e) {
            sendError(exchange, 500, "错误: " + e.getMessage());
        }
    }
    
    private void handleProgramInfo(HttpExchange exchange) throws IOException {
        try {
            JsonObject response = new JsonObject();
            response.addProperty("name", currentProgram.getName());
            response.addProperty("executable_path", currentProgram.getExecutablePath());
            response.addProperty("language", currentProgram.getLanguage().getLanguageID().toString());
            response.addProperty("compiler", currentProgram.getCompiler());
            response.addProperty("image_base", currentProgram.getImageBase().toString());
            
            Memory memory = currentProgram.getMemory();
            response.addProperty("memory_size", memory.getSize());
            
            FunctionManager funcManager = currentProgram.getFunctionManager();
            response.addProperty("function_count", funcManager.getFunctionCount());
            
            sendResponse(exchange, 200, response);
            
        } catch (Exception e) {
            sendError(exchange, 500, "错误: " + e.getMessage());
        }
    }
    private void handleAsm(HttpExchange exchange) throws IOException {
        Program program = currentProgram;
        JsonObject request = parseRequest(exchange);

        String addrStr = request.get("address").getAsString();
        Address addr = program.getAddressFactory().getAddress(addrStr);
        Function function = program.getFunctionManager().getFunctionAt(addr);

        JsonObject response = new JsonObject();
        if (function == null) {
            response.addProperty("error", "Function not found at " + addrStr);
        } else {
            response.addProperty("function", function.getName());
            response.addProperty("address", function.getEntryPoint().toString());

            JsonArray instructions = new JsonArray();
            InstructionIterator iter = program.getListing().getInstructions(function.getBody(), true);
            while (iter.hasNext()) {
                Instruction instr = iter.next();
                JsonObject instrObj = new JsonObject();
                instrObj.addProperty("address", instr.getAddress().toString());
                instrObj.addProperty("mnemonic", instr.getMnemonicString());

                // 使用 getDefaultOperandRepresentation(i) 获取完整字符串
                JsonArray ops = new JsonArray();
                int opCount = instr.getNumOperands();
                for (int i = 0; i < opCount; i++) {
                    String opStr = instr.getDefaultOperandRepresentation(i);
                    ops.add(opStr);
                }
                instrObj.add("operands", ops);

                instructions.add(instrObj);
            }
            response.add("instructions", instructions);
            response.addProperty("count", instructions.size());
        }

        sendResponse(exchange, 200, response);
    }

    private void handleHealth(HttpExchange exchange) throws IOException {
        JsonObject response = new JsonObject();
        response.addProperty("status", "ok");
        response.addProperty("program_loaded", currentProgram != null);
        sendResponse(exchange, 200, response);
    }

    // ========== 内存相关API ==========

    private void handleMemory(HttpExchange exchange) throws IOException {
        if (!"POST".equals(exchange.getRequestMethod())) {
            sendError(exchange, 405, "不允许的方法");
            return;
        }
        
        try {
            JsonObject request = parseRequest(exchange);
            String addressStr = request.get("address").getAsString();
            int length = request.get("length").getAsInt();
            
            // 验证长度范围
            if (length < 1 || length > 1024) {
                sendError(exchange, 400, "长度必须在1-1024字节之间");
                return;
            }
            
            Address addr = getAddress(addressStr);
            Memory memory = currentProgram.getMemory();
            
            // 检查地址是否在有效内存范围内
            if (!memory.contains(addr)) {
                sendError(exchange, 400, "地址不在有效内存范围内");
                return;
            }
            
            // 检查请求的内存范围是否有效
            Address endAddr = addr.add(length - 1);
            if (!memory.contains(endAddr)) {
                sendError(exchange, 400, "请求的内存范围超出有效边界");
                return;
            }
            
            // 读取内存数据
            byte[] buffer = new byte[length];
            memory.getBytes(addr, buffer);
            
            // 将字节数据转换为十六进制字符串
            StringBuilder hexBuilder = new StringBuilder();
            for (byte b : buffer) {
                hexBuilder.append(String.format("%02X", b));
            }
            
            // 生成ASCII表示
            StringBuilder asciiBuilder = new StringBuilder();
            for (byte b : buffer) {
                asciiBuilder.append(b >= 32 && b <= 126 ? (char) b : '.');
            }
            
            JsonObject response = new JsonObject();
            response.addProperty("address", addressStr);
            response.addProperty("length", length);
            response.addProperty("hex", hexBuilder.toString());
            response.addProperty("ascii", asciiBuilder.toString());
            
            sendResponse(exchange, 200, response);
            
        } catch (Exception e) {
            sendError(exchange, 500, "错误: " + e.getMessage());
        }
    }
    
    private void handleMemoryRegions(HttpExchange exchange) throws IOException {
        if (!"GET".equals(exchange.getRequestMethod())) {
            sendError(exchange, 405, "不允许的方法");
            return;
        }
        
        try {
            JsonArray regions = new JsonArray();
            Memory memory = currentProgram.getMemory();
            
            // 获取所有内存区域
            for (MemoryBlock block : memory.getBlocks()) {
                JsonObject region = new JsonObject();
                region.addProperty("name", block.getName());
                region.addProperty("start_address", block.getStart().toString());
                region.addProperty("end_address", block.getEnd().toString());
                region.addProperty("size", block.getSize());
                region.addProperty("read", block.isRead());
                region.addProperty("write", block.isWrite());
                region.addProperty("execute", block.isExecute());
                region.addProperty("initialized", block.isInitialized());
                
                regions.add(region);
            }
            
            JsonObject response = new JsonObject();
            response.add("regions", regions);
            response.addProperty("count", regions.size());
            
            sendResponse(exchange, 200, response);
            
        } catch (Exception e) {
            sendError(exchange, 500, "错误: " + e.getMessage());
        }
    }
    
    private void handleMemorySearch(HttpExchange exchange) throws IOException {
        if (!"POST".equals(exchange.getRequestMethod())) {
            sendError(exchange, 405, "不允许的方法");
            return;
        }
        
        try {
            JsonObject request = parseRequest(exchange);
            String pattern = request.get("pattern").getAsString();
            String startAddressStr = request.has("start_address") ? request.get("start_address").getAsString() : null;
            String endAddressStr = request.has("end_address") ? request.get("end_address").getAsString() : null;
            int limit = request.has("limit") ? request.get("limit").getAsInt() : 100;
            
            // 转换十六进制模式为字节数组
            byte[] searchPattern = hexStringToByteArray(pattern);
            if (searchPattern.length == 0) {
                sendError(exchange, 400, "无效的搜索模式");
                return;
            }
            
            Memory memory = currentProgram.getMemory();
            AddressSetView searchRange;
            
            // 设置搜索范围
            if (startAddressStr != null && endAddressStr != null) {
                Address start = getAddress(startAddressStr);
                Address end = getAddress(endAddressStr);
                searchRange = new AddressSet(start, end);
            } else {
                searchRange = memory.getLoadedAndInitializedAddressSet();
            }
            
            // 搜索内存
            JsonArray results = new JsonArray();
            int count = 0;
            
            AddressIterator addressIterator = searchRange.getAddresses(true);
            while (addressIterator.hasNext() && count < limit) {
                Address addr = addressIterator.next();
                
                // 检查是否有足够空间来比较模式
                if (addr.add(searchPattern.length - 1).compareTo(memory.getMaxAddress()) > 0) {
                    continue;
                }
                
                // 读取当前地址的字节
                byte[] buffer = new byte[searchPattern.length];
                try {
                    memory.getBytes(addr, buffer);
                    
                    // 比较字节数组
                    boolean match = true;
                    for (int i = 0; i < searchPattern.length; i++) {
                        if (buffer[i] != searchPattern[i]) {
                            match = false;
                            break;
                        }
                    }
                    
                    if (match) {
                        JsonObject result = new JsonObject();
                        result.addProperty("address", addr.toString());
                        results.add(result);
                        count++;
                    }
                } catch (Exception e) {
                    // 跳过不可访问的内存
                    continue;
                }
            }
            
            JsonObject response = new JsonObject();
            response.add("results", results);
            response.addProperty("count", results.size());
            
            sendResponse(exchange, 200, response);
            
        } catch (Exception e) {
            sendError(exchange, 500, "错误: " + e.getMessage());
        }
    }
    
    private void handleMemoryWrite(HttpExchange exchange) throws IOException {
        if (!"POST".equals(exchange.getRequestMethod())) {
            sendError(exchange, 405, "不允许的方法");
            return;
        }
        
        try {
            JsonObject request = parseRequest(exchange);
            String addressStr = request.get("address").getAsString();
            String data = request.get("data").getAsString();
            String dataType = request.has("data_type") ? request.get("data_type").getAsString() : "hex";
            
            Address addr = getAddress(addressStr);
            Memory memory = currentProgram.getMemory();
            
            // 检查地址是否可写
            MemoryBlock block = memory.getBlock(addr);
            if (block == null || !block.isWrite()) {
                sendError(exchange, 403, "地址所在的内存区域不可写");
                return;
            }
            
            // 转换数据格式
            byte[] bytes;
            if ("hex".equalsIgnoreCase(dataType)) {
                bytes = hexStringToByteArray(data);
            } else if ("ascii".equalsIgnoreCase(dataType)) {
                bytes = data.getBytes("UTF-8");
            } else {
                sendError(exchange, 400, "不支持的数据类型: " + dataType);
                return;
            }
            
            // 检查是否有足够的空间
            Address endAddr = addr.add(bytes.length - 1);
            if (!memory.contains(endAddr)) {
                sendError(exchange, 400, "写入的数据将超出有效内存边界");
                return;
            }
            
            // 开始事务并写入内存
            int txId = currentProgram.startTransaction("写入内存数据");
            try {
                memory.setBytes(addr, bytes);
                currentProgram.endTransaction(txId, true);
            } catch (Exception e) {
                currentProgram.endTransaction(txId, false);
                sendError(exchange, 500, "写入内存失败: " + e.getMessage());
                return;
            }
            
            JsonObject response = new JsonObject();
            response.addProperty("success", true);
            response.addProperty("address", addressStr);
            response.addProperty("written_bytes", bytes.length);
            
            sendResponse(exchange, 200, response);
            
        } catch (Exception e) {
            sendError(exchange, 500, "错误: " + e.getMessage());
        }
    }
    
    // ========== 辅助方法 ==========
    
    private JsonObject parseRequest(HttpExchange exchange) throws IOException {
        InputStreamReader isr = new InputStreamReader(exchange.getRequestBody(), "UTF-8");
        return gson.fromJson(isr, JsonObject.class);
    }
    
    private Address getAddress(String addressStr) {
        return currentProgram.getAddressFactory().getAddress(addressStr);
    }
    
    private void sendResponse(HttpExchange exchange, int statusCode, JsonObject data) throws IOException {
        String response = gson.toJson(data);
        byte[] bytes = response.getBytes("UTF-8");
        
        exchange.getResponseHeaders().set("Content-Type", "application/json; charset=UTF-8");
        exchange.getResponseHeaders().set("Access-Control-Allow-Origin", "*");
        exchange.sendResponseHeaders(statusCode, bytes.length);
        
        OutputStream os = exchange.getResponseBody();
        os.write(bytes);
        os.close();
    }
    
    private void sendError(HttpExchange exchange, int statusCode, String message) throws IOException {
        JsonObject error = new JsonObject();
        error.addProperty("error", message);
        sendResponse(exchange, statusCode, error);
    }
    
    // 将十六进制字符串转换为字节数组
    private byte[] hexStringToByteArray(String s) {
        int len = s.length();
        byte[] data = new byte[len / 2];
        for (int i = 0; i < len; i += 2) {
            data[i / 2] = (byte) ((Character.digit(s.charAt(i), 16) << 4)
                                 + Character.digit(s.charAt(i+1), 16));
        }
        return data;
    }
    
    private void stopServer() {
        if (server != null) {
            server.stop(0);
            println("服务器已停止。");
        }
        if (decompiler != null) {
            decompiler.dispose();
        }
    }
}